﻿using org.apache.zookeeper.data;
using System.Text;

namespace JESAI.DistributedLock.ZooKeeper;

using JESAI.DistributedLock.Internal;
using org.apache.zookeeper;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;

/// <summary>
/// 为ZooKeeper同步算法提供一个通用结构。该结构如下：
/// * 创建一个顺序临时节点
/// * 循环
///     * 将我们节点的序列号与其兄弟节点进行比较，以查看我们是否已获取
///     * 如果我们尚未获取，等待一个监听器告诉我们某些事情已发生变化
/// </summary>
internal class ZooKeeperSynchronizationHelper
{
    public static readonly IReadOnlyList<byte> AcquiredMarker = Encoding.UTF8.GetBytes("ACQUIRED");

    private readonly ZooKeeperConnectionInfo _connectionInfo;
    private readonly IReadOnlyList<ACL> _acl;
    private readonly bool _assumePathExists, _setAcquiredMarker;

    public ZooKeeperSynchronizationHelper(
        ZooKeeperPath path,
        bool assumePathExists,
        string connectionString,
        Action<ZooKeeperDistributedSynchronizationOptionsBuilder>? optionsBuilder,
        bool setAcquiredMarker = false)
    {
        this.Path = path;
        this._assumePathExists = assumePathExists;
        var options = ZooKeeperDistributedSynchronizationOptionsBuilder.GetOptions(optionsBuilder);
        this._connectionInfo = new ZooKeeperConnectionInfo(
            connectionString ?? throw new ArgumentNullException(nameof(connectionString)),
            ConnectTimeout: options.ConnectTimeout,
            SessionTimeout: options.SessionTimeout,
            AuthInfo: options.AuthInfo
        );
        this._acl = options.Acl;
        this._setAcquiredMarker = setAcquiredMarker;
    }

    public ZooKeeperPath Path { get; }

    public async Task<ZooKeeperNodeHandle?> TryAcquireAsync(
        Func<State, bool> hasAcquired,
        Func<ZooKeeper, State, Watcher, Task<bool>> waitAsync,
        TimeoutValue timeout,
        CancellationToken cancellationToken,
        string nodePrefix,
        string? alternateNodePrefix = null)
    {
        var acquired = false;
        var ephemeralNodeLost = false;
        ZooKeeperConnection? connection = null;
        string? ephemeralNodePath = null;
        var timeoutSource = new CancellationTokenSource(timeout.TimeSpan);

        try
        {
            connection = await ZooKeeperConnection.DefaultPool.ConnectAsync(this._connectionInfo, cancellationToken).ConfigureAwait(false);

            // 创建一条路径，既表示我们对synch对象的持有，也表示我们在等待它的队列中的位置
            ephemeralNodePath = await connection.CreateEphemeralSequentialNode(this.Path, nodePrefix, this._acl, ensureDirectoryExists: !this._assumePathExists).ConfigureAwait(false);

            while (true)
            {
                // 获取节点的子节点，按前缀过滤，并按年龄排序；我们将用它来确定我们在列表中的位置
                var children = await connection.ZooKeeper.getChildrenAsync(this.Path.ToString()).ConfigureAwait(false);
                var sortedChildren = await ZooKeeperSequentialPathHelper.FilterAndSortAsync(
                    parentNode: this.Path.ToString(),
                    childrenNames: children.Children,
                    getNodeCreationTimeAsync: connection.GetNodeCreationTimeAsync,
                    prefix: nodePrefix,
                    alternatePrefix: alternateNodePrefix
                ).ConfigureAwait(false);

                // 卫生检查；如果其他人从我们下面删除它，可能会发生这种情况。我们必须先检查一下，因为它涵盖了空的收集箱
                if (!sortedChildren.Any(t => t.Path == ephemeralNodePath))
                {
                    ephemeralNodeLost = true;
                    throw new InvalidOperationException($"Node '{ephemeralNodePath}' was created, but no longer exists");
                }

                // 看看我们是否已经收购
                var state = new State(ephemeralNodePath, sortedChildren);
                if (hasAcquired(state))
                {
                    if (this._setAcquiredMarker)
                    {
                        await connection.ZooKeeper.setDataAsync(ephemeralNodePath, AcquiredMarker.ToArray()).ConfigureAwait(false);
                    }
                    acquired = true;
                    return new ZooKeeperNodeHandle(connection, ephemeralNodePath, shouldDeleteParent: !this._assumePathExists);
                }

                // 等待变化
                var waitCompletionSource = new TaskCompletionSource<bool>();
                using var timeoutRegistration = timeoutSource.Token.Register(state => ((TaskCompletionSource<bool>)state).TrySetResult(false), waitCompletionSource);
                using var cancellationRegistration = cancellationToken.Register(state => ((TaskCompletionSource<bool>)state).TrySetCanceled(), waitCompletionSource);
                // 这是必要的，因为如果连接中断且无法恢复，我们将永远不会收到会话过期通知。
                using var connectionLostRegistration = connection.ConnectionLostToken.Register(
                    state => ((TaskCompletionSource<bool>)state).TrySetException(new InvalidOperationException("Lost connection to ZooKeeper")),
                    state: waitCompletionSource
                );
                if (!waitCompletionSource.Task.IsCompleted
                    && await waitAsync(connection.ZooKeeper, state, new WaitCompletionSourceWatcher(waitCompletionSource)).ConfigureAwait(false))
                {
                    waitCompletionSource.TrySetResult(true);
                }

                if (!await waitCompletionSource.Task.ConfigureAwait(false))
                {
                    return null; // 等待超时
                }
            }
        }
        finally
        {
            timeoutSource.Dispose();

            // 如果我们未能获取，请清理连接/节点路径
            if (!acquired)
            {
                await this.CleanUpOnFailureAsync(connection, ephemeralNodePath, ephemeralNodeLost).ConfigureAwait(false);
            }
        }
    }

    private async Task CleanUpOnFailureAsync(ZooKeeperConnection? connection, string? ephemeralNodePath, bool ephemeralNodeLost)
    {
        if (connection == null) { return; }

        try
        {
            if (ephemeralNodePath != null)
            {
                if (!ephemeralNodeLost)
                {
                    await connection.ZooKeeper.deleteAsync(ephemeralNodePath).ConfigureAwait(false);
                }
                if (!this._assumePathExists)
                {
                    // 如果父节点应该被清理，尝试这样做。这个尝试几乎肯定会失败，因为
                    // 其他人持有锁。然而，我们可能遇到了一个竞争条件，即另一个持有者
                    // 在我们获取失败后立即释放了锁，而我们的临时节点阻止了他们删除。因此，我们
                    // 忽略这个删除操作，以覆盖这种情况，而不会减慢我们的速度

                    _ = connection.ZooKeeper.deleteAsync(this.Path.ToString());
                }
            }
        }
        finally
        {
            connection.Dispose();
        }
    }

    private sealed class WaitCompletionSourceWatcher : Watcher
    {
        private readonly TaskCompletionSource<bool> _waitCompletionSource;

        public WaitCompletionSourceWatcher(TaskCompletionSource<bool> waitCompletionSource)
        {
            this._waitCompletionSource = waitCompletionSource;
        }

        public override Task process(WatchedEvent @event)
        {
            // 只关心相关状态事件；ConnectionLostToken为我们处理其他状态
            if (@event.getState() == Event.KeeperState.SyncConnected)
            {
                this._waitCompletionSource.TrySetResult(true);
            }

            return Task.CompletedTask;
        }
    }

    public record State(string EphemeralNodePath, (string Path, int SequenceNumber, string Prefix)[] SortedChildren);
}
